import os
import re
import shutil
import zipfile
import docx
import requests  # 使用requests库调用硅基流动API
import tkinter as tk
from tkinter import filedialog
import argparse

# 配置
OUTPUT_DIR = 'output'
EXTRACT_DIR = 'extracted_files'  # 专门建个文件夹放解压内容
API_KEY = 'sk-kzqzzhnyszoknmvmlkeysyzawazylgatpkeeqlivnfirbqgo'  # 用户提供的硅基流动API密钥
USER_RULES = []  # 存储用户定义的规则

# 硅基流动 API配置
SILICONFLOW_API_URL = 'https://api.siliconflow.cn/v1/chat/completions'
SILICONFLOW_MODEL = 'THUDM/GLM-4-9B-0414'  # 经测试有效的模型名称

def clean_filename(filename):
    """清理文件名中的无效字符"""
    # 移除控制字符（回车、换行、制表符等）
    filename = re.sub(r'[\r\n\t]', '', filename)
    # 替换Windows文件系统不允许的字符
    invalid_chars = '<>:"|?*'
    for char in invalid_chars:
        filename = filename.replace(char, '_')
    # 移除多余的下划线
    filename = re.sub(r'_+', '_', filename)
    # 确保文件名不为空
    if not filename.strip():
        filename = "unnamed_file"
    return filename

def extract_zip(zip_path, extract_to):
    """解压ZIP文件，支持中文路径和特殊字符处理"""
    try:
        # 确保解压目录存在
        os.makedirs(extract_to, exist_ok=True)
        
        with zipfile.ZipFile(zip_path, 'r') as zip_ref:
            # 修复中文文件名乱码问题并清理无效字符
            for info in zip_ref.filelist:
                # 尝试用GBK解码文件名（Windows下常见编码）
                try:
                    correct_filename = info.filename.encode('cp437').decode('gbk')
                except (UnicodeDecodeError, UnicodeEncodeError):
                    # 如果GBK失败，尝试UTF-8
                    try:
                        correct_filename = info.filename.encode('cp437').decode('utf-8')
                    except (UnicodeDecodeError, UnicodeEncodeError):
                        # 都不行就用原始文件名
                        correct_filename = info.filename
                
                # 清理文件名中的无效字符
                correct_filename = clean_filename(correct_filename)
                
                # 更新文件名并解压
                info.filename = correct_filename
                zip_ref.extract(info, extract_to)
                
        print(f"已解压到: {extract_to}")
    except Exception as e:
        print(f"解压失败: {e}")
        raise

def read_docx(file_path):
    """读取Word文档内容"""
    try:
        doc = docx.Document(file_path)
        return '\n'.join([para.text for para in doc.paragraphs])
    except Exception as e:
        print(f"读取文件失败 {file_path}: {e}")
        return ""

def get_user_rules():
    """获取用户定义的规则"""
    global USER_RULES
    
    # 默认规则（直接使用默认规则，无需用户输入）
    rules = ["描述了一个项目", "包含具体的技术实现", "有明确的目标或需求"]
    print(f"已设置默认规则: {'; '.join(rules)}")
    
    USER_RULES = rules
    return rules

def is_target_project(text):
    """用硅基流动DeepSeek-R1模型根据用户规则判断文档是否符合要求"""
    if not text.strip():
        return False
    
    if not USER_RULES:
        # 如果没有设置规则，使用默认规则
        rules = ["描述了一个项目", "包含具体的技术实现", "有明确的目标或需求"]
    else:
        rules = USER_RULES
    
    try:
        # 设置请求头和数据
        headers = {
            "Authorization": f"Bearer {API_KEY}",
            "Content-Type": "application/json"
        }
        
        # 构建规则提示
        rules_text = "\n".join([f"{i+1}. {rule}" for i, rule in enumerate(rules)])
        system_prompt = f"""我将给你一系列规则和一段文本，
请判断这段文本是否符合以下规则中至少一条的定义，只需回复'是'或'否'，不要添加其他内容。
规则：
{rules_text}"""
        
        # 发送请求
        data = {
            "model": SILICONFLOW_MODEL,
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": text[:1000]}  # 取前1000字符，提供更多上下文
            ],
            "max_tokens": 2,
            "temperature": 0.0
        }
        
        print(f"正在调用硅基流动API，使用模型: {SILICONFLOW_MODEL}")
        response = requests.post(SILICONFLOW_API_URL, headers=headers, json=data)
        
        # 处理响应
        if response.status_code == 200:
            result = response.json()
            content = result['choices'][0]['message']['content'].strip().lower()
            print(f"模型返回: {content}")
            return '是' in content
        else:
            print(f"API调用失败: 状态码{response.status_code}, 错误信息:{response.text}")
            return False
    except Exception as e:
        print(f"API调用失败: {e}")
        return False

def process_files(zip_path):
    """主处理流程"""
    # 创建必要的目录
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    
    # 解压ZIP到指定目录
    extract_zip(zip_path, EXTRACT_DIR)
    
    # 记录处理的文件数
    processed_count = 0
    
    # 在解压目录中处理每个Word文档
    for root, _, files in os.walk(EXTRACT_DIR):
        for file in files:
            file_lower = file.lower()
            if file_lower.endswith('.docx') or file_lower.endswith('.doc'):
                file_path = os.path.join(root, file)
                print(f"正在处理: {file}")
                
                content = read_docx(file_path) if file_lower.endswith('.docx') else "暂不支持.doc文件内容读取"
                
                # 对于.doc文件，直接复制到输出目录，不做AI判断
                if file_lower.endswith('.doc'):
                    print(f"⚠ 发现.doc文件，直接复制: {file}")
                    dest_path = os.path.join(OUTPUT_DIR, file)
                    
                    # 如果目标文件已存在，加个序号避免覆盖
                    counter = 1
                    base_name, ext = os.path.splitext(file)
                    while os.path.exists(dest_path):
                        dest_path = os.path.join(OUTPUT_DIR, f"{base_name}_{counter}{ext}")
                        counter += 1
                    
                    shutil.copy2(file_path, dest_path)
                    processed_count += 1
                    print(f"✓ 已复制.doc文件: {os.path.basename(dest_path)}")
                elif content and is_target_project(content):
                    # 处理.docx文件
                    dest_path = os.path.join(OUTPUT_DIR, file)
                    
                    # 如果目标文件已存在，加个序号避免覆盖
                    counter = 1
                    base_name, ext = os.path.splitext(file)
                    while os.path.exists(dest_path):
                        dest_path = os.path.join(OUTPUT_DIR, f"{base_name}_{counter}{ext}")
                        counter += 1
                    
                    shutil.copy2(file_path, dest_path)
                    processed_count += 1
                    print(f"✓ 已复制意向项目文件: {os.path.basename(dest_path)}")
    
    print(f"处理完成！共找到 {processed_count} 个意向项目文件，保存在 {OUTPUT_DIR} 目录下")
    print(f"注意：文件名显示乱码是ZIP解压编码问题，文件内容不受影响")

def select_zip_file():
    """打开文件选择对话框选择ZIP文件"""
    root = tk.Tk()
    root.withdraw()  # 隐藏主窗口
    
    file_path = filedialog.askopenfilename(
        title="选择ZIP文件",
        filetypes=[("ZIP files", "*.zip")]
    )
    return file_path

def main():
    # 创建命令行参数解析器
    parser = argparse.ArgumentParser(description='处理ZIP文件中的Word文档并筛选意向项目')
    parser.add_argument('zip_path', nargs='?', help='ZIP文件路径')
    
    args = parser.parse_args()
    
    # 如果未提供ZIP路径，使用文件选择对话框
    if not args.zip_path:
        args.zip_path = select_zip_file()
        if not args.zip_path:
            print("未选择ZIP文件，程序退出。")
            return
    
    # 检查ZIP文件是否存在
    if not os.path.exists(args.zip_path):
        print(f"错误: ZIP文件 '{args.zip_path}' 不存在！")
        return
    
    # 检查是否为ZIP文件
    if not args.zip_path.lower().endswith('.zip'):
        print(f"错误: '{args.zip_path}' 不是一个ZIP文件！")
        return
    
    # 检查API密钥是否为空
    global API_KEY
    if not API_KEY or API_KEY == 'sk-kzqzzhnyszoknmvmlkeysyzawazylgatpkeeqlivnfirbqgo':
        # 询问用户是否需要修改API密钥
        use_default = input("是否使用默认API密钥？(Y/n): ").strip().lower()
        if use_default != 'y' and use_default != '':
            user_api_key = input("请输入您的硅基流动API密钥: ").strip()
            if user_api_key:
                API_KEY = user_api_key
            else:
                print("未输入新密钥，继续使用默认密钥。")
        else:
            print("正在使用默认API密钥。")
    
    # 获取用户定义的规则
    get_user_rules()
    
    try:
        # 执行主处理流程
        process_files(args.zip_path)
    except Exception as e:
        print(f"处理过程中发生错误: {e}")

if __name__ == "__main__":
    main()